#!/usr/bin/python
'''
print ace file from annotated ham file.
'''
import sys
import ham
import h5py
import logging
import numpy as np
import fasta_read

logging.basicConfig(level=logging.INFO, format='[%(levelname)s] (%(threadName)-10s) %(message)s', )
logging.basicConfig(level=logging.DEBUG, format='[%(levelname)s] (%(threadName)-10s) %(message)s', )

# parameters.
ham_file1 = sys.argv[1]
ham_file2 = sys.argv[2]
ant_file = sys.argv[3]
ref_file = sys.argv[4]
ace_file = sys.argv[5]
mates_file = sys.argv[6]

dummy1 = "ACGTTACGTTACGTTACGTTACGTTACGTTACGTTACGTTACGTTACGTT"
dummy2 = "GCCTAGCCTAGCCTAGCCTAGCCTAGCCTAGCCTAGCCTAGCCTAGCCTA"

# read in reference.
logging.info("reading reference.")
ref = fasta_read.read_fasta_dict(ref_file)

# open and load annotation.
#logging.info("loading annotations.")
#ant_obj = ham.HamAnnotate(ant_file, read=1)
#ant_arr1, ant_arr2 = ant_obj.load_annotation()
#ant_arr1 = ant_obj.h5_ant_in['ant_a'][0:10000]
#ant_arr2 = ant_obj.h5_ant_in['ant_b'][0:10000]


# open ham files.
logging.info("loading ham file columns.")
ham_obj1 = ham.HamFile(ham_file1, read=1)
ham_obj2 = ham.HamFile(ham_file2, read=1)
#ham_arr1 = ham_obj1.data[:, 'FLAG', 'RNAME', 'POS']
#ham_arr2 = ham_obj2.data[:, 'FLAG', 'RNAME', 'POS']
ham_arr1 = ham_obj1.data[:, 'QNAME','FLAG', 'RNAME', 'POS']
ham_arr2 = ham_obj2.data[:, 'QNAME','FLAG', 'RNAME', 'POS']

# print loaded the following.
print "loaded the follwoing", ham_arr1.size, ham_arr2.size

# start mates file.
fout = open(mates_file, "wb")
fout.write("library\tA\t645\t2145\n")

# find double unique partners.
logging.info("building ace file.")
ace = {}
cnt = 0
namecnt = 0
for i in range(ham_arr1.size):
	'''
	# look for uniquely mapped.
	if ant_arr1[i]['TYPE'][0] != 1: continue
	
	# look for noted partner.
	if ant_arr1[i]['MIDX'][0] == -1: continue
	j = ant_arr1[i]['MIDX'][0]
	
	if j > ham_arr2.size: continue
	
	# ensure partner is unique.
	if ant_arr2[j]['TYPE'][0] != 1: continue

	# ensure they match up.
	if ant_arr2[j]['MIDX'][0] != i:
		logging.error("bad matching pair: %i %i" % (i, j))
		sys.exit(1)
	'''
	
	# get data.
	rname1 = ham_arr1[i]['RNAME'][0]
	rname2 = ham_arr2[i]['RNAME'][0]
#	qname1 = ham_arr1[i]['QNAME'][0]
#	qname2 = ham_arr2[j]['QNAME'][0]

	if rname1 == '' or rname2 == '': continue

	if rname1 == "*":
		continue
		print "1"
		print ham_arr1[i]
		
		print "2"
		print ham_arr2[i]
		sys.exit()
		
	if rname2 == "*":
		continue
		print "2"
		print ham_arr2[i]
		
		print "1"
		print ham_arr1[i]
		sys.exit()		
	
	qname1 = "name1_%i" % namecnt
	qname2 = "name2_%i" % namecnt
	namecnt += 1
	
	pos1 = ham_arr1[i]['POS']
	pos2 = ham_arr2[i]['POS']
	#seq1 = ham_arr1[i]['SEQ']
	#seq2 = ham_arr2[j]['SEQ']
	seq1 = dummy1
	seq2 = dummy2
	szseq1 = len(seq1)
	szseq2 = len(seq2)
	
	if ham_arr1[i]['FLAG'] == 0: orien1 = "U"
	else: orien1 = "C"
	if ham_arr2[i]['FLAG'] == 0: orien2 = "U"
	else: orien2 = "C"
	
	
	# write ace txt.
	AF1 = "AF %s %s %i\n" % (qname1, orien1, pos1)
	RD1 = "RD %s %i 0 0\n" % (qname1, szseq1)
	RD1 += seq1 + "\n\n"
	RD1 += "QA 1 %s 1 %s\n" % (szseq1, szseq1)
	RD1 += "DS CHROMAT_FILE: ENV1180.x1 PHD_FILE: ENV1180.x1.phd.1 TIME: Wed Sep 15 15:46:47 1999\n\n"
	
	AF2 = "AF %s %s %i\n" % (qname2, orien2, pos2)
	RD2 = "RD %s %i 0 0\n" % (qname2, szseq2)
	RD2 += seq1 + "\n\n"
	RD2 += "QA 1 %s 1 %s\n" % (szseq2, szseq2)
	RD2 += "DS CHROMAT_FILE: ENV1180.x1 PHD_FILE: ENV1180.x1.phd.1 TIME: Wed Sep 15 15:46:47 1999\n\n"
	
	# save to ace structure.
	if rname1 not in ace:
		ace[rname1] = ([], [], [], [])
	ace[rname1][0].append(AF1)
	ace[rname1][1].append(RD1)
#	ace[rname1][2].append(QA1)
#	ace[rname1][3].append(DS1)
	
	if rname2 not in ace:
		ace[rname2] = ([], [], [], [])
	ace[rname2][0].append(AF2)
	ace[rname2][1].append(RD2)
#	ace[rname2][2].append(QA2)
#	ace[rname2][3].append(DS2)

	cnt += 2
	
	# write pair to mates.
	fout.write("%s\t%s\tA\n" % (qname1, qname2))
fout.close()

# print out ace file.
logging.info("writing ace file.")

fout = open(ace_file, "wb")
fout.write("AS %i %i\n\n" % (len(ref), cnt))

for ctg in ace:
	# build ctg info.
	cseq = ref[ctg]
	szcseq = len(cseq)
	chits = ace[ctg]
	txt = "CO %s %i %i 0 U\n" % (ctg, szcseq, len(chits))
	txt += cseq + "\n\n"
	txt += "BQ\n"
	txt += "20 " * szcseq + "\n\n"
	
	# write contig info.
	fout.write(txt)
	
	# write read info.
	for e in chits[0]:
		fout.write(e)
	fout.write("\n")
		
	for e in chits[1]:
		fout.write(e)
		

# close h5 files.
fout.close()
#ant_obj.close()
ham_obj1.close()
ham_obj2.close()
